package scan

import (
	"Suscan/models"
	"Suscan/pkg/masscan"
	"fmt"
	"net"
	"regexp"
	"strings"
	"sync"
	"time"
)

type MassScanRes struct {
	Ip       string
	Port     string
	Protocol string
}

func InitMasscan() {
	start := time.Now()
	massRes := MasscanStart()

	costTime := time.Since(start)
	data := make(map[string]interface{})
	data["taskid"] = 0
	data["task_name"] = "masscan"
	data["task_type"] = "masscan"
	data["all_num"] = len(massRes)
	data["succes_num"] = len(massRes)
	data["run_time"] = fmt.Sprintf("%s", costTime)
	models.AddLog(data)

	wg := &sync.WaitGroup{}
	thread := 2
	taskChan := make(chan MassScanRes, 50*2)

	for i := 0; i < thread; i++ {
		go NmapScan(taskChan, wg)
	}

	for _, task := range massRes {
		wg.Add(1)
		taskChan <- task
	}

	close(taskChan)
	wg.Wait()
}

func MasscanStart() (massRes []MassScanRes) {
	rate := "3000"
	massRes = make([]MassScanRes, 0)
	result := models.GetAllAsset()

	var ips []string

	for _, r := range result {
		//ip与url转为ip
		address, err := net.LookupHost(r.Target)
		//如果是target格式是ip，直接加入扫描列表
		if err == nil{
			m, _ := regexp.MatchString("^(25[0-5]|2[0-4]\\d|[0-1]\\d{2}|[1-9]?\\d)\\.(25[0-5]|2[0-4]\\d|[0-1]\\d{2}|[1-9]?\\d)\\.(25[0-5]|2[0-4]\\d|[0-1]\\d{2}|[1-9]?\\d)\\.(25[0-5]|2[0-4]\\d|[0-1]\\d{2}|[1-9]?\\d)$", address[0])
			if m {
				ip := address[0]
				ips = append(ips, ip)
				ok, id := models.Existurl(r.Target,ip)
				if ok {
					//更新update时间
					dataUpdate := make(map[string]interface{})
					nowTime := time.Now().Format("2006-01-02 15:04:05")
					dataUpdate["updated_time"] =nowTime
					models.EditUrl(id,dataUpdate)
				}else {
					data := make(map[string]interface{})
					data["url"] = r.Target
					data["ip"] = ip
					models.AddUrl(data)
				}
			}
		}
	}

	var ipsresult string

	ipsresult = strings.Join(ips,",")

	m := masscan.New()
	m.SetSystemPath("masscan")
	m.SetRate(rate)

	args := []string{
		ipsresult,
		//"-p", "80,7799,8080,443",
		"-p", "1-65535",
		//"--exclude", ipsNoScan.String(),
	}
	m.SetArgs(args...)

	err := m.Run()
	if err != nil {
		fmt.Println("mass run  err : ", err)
	}

	results, err := m.Parse()
	if err != nil {
		fmt.Println(err)
	}
	fmt.Printf("Masscan find %d ips \n", len(results))

	for _, result := range results {

		massResTmp := MassScanRes{Ip: result.Address.Addr, Port: result.Ports[0].Portid}
		massRes = append(massRes, massResTmp)
	}

	fmt.Println("mass dup is ", massRes)

	return
}

// 读取mass扫描结果，跟当前数据库对比。如果不存在，就namp识别端口入库

func NmapScan(taskChan chan MassScanRes, wg *sync.WaitGroup) {

	data := make(map[string]interface{})
	dataUpdate := make(map[string]interface{})
	for target := range taskChan {
		defer func() {
			if err := recover(); err != nil {
				wg.Done()
			}
		}()

		data["ip"] = target.Ip
		data["port"] = target.Port
		//扫描结果入库前对比
		//对比ip和端口，返回bool+ iplist.ID
		ok, id := models.ExistIplist(target.Ip, target.Port)

		//对比数据库中ip是否存在，如果存在的话，url字段也存在。返回

		_,url:=models.Searchurl(target.Ip)

		if ok {
			fmt.Println(target.Ip, target.Port, "ip,port存在数据库中 更新时间")
			nowTime := time.Now().Format("2006-01-02 15:04:05")
			dataUpdate["updated_time"] = nowTime
			models.EditIplist(id, dataUpdate)
			wg.Done()
		} else {
			fmt.Println(target.Ip, target.Port, "不存在数据库中,nmap识别端口指纹，入库")
			IpNmap(target.Ip, target.Port, url)
			wg.Done()
		}
	}
}
